package com.y.report.util;

import com.alibaba.fastjson.JSON;
import org.apache.commons.lang.StringUtils;
import org.apache.http.*;
import org.apache.http.client.methods.*;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.DefaultRedirectStrategy;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.protocol.HttpContext;
import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.net.HttpURLConnection;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.*;

public class DorisStreamLoad implements Serializable {

    private static final Logger log = LoggerFactory.getLogger(DorisStreamLoad.class);
    //连接地址，这里使用的是连接FE
    private static String loadUrlPattern = "http://%s/api/%s/%s/_stream_load";
    //fe ip地址
    private String hostPort;
    //数据库
    private String db;
    //要导入的数据表名
    private String tbl;
    //用户名
    private String user;
    //密码
    private String passwd;
    private String authEncoding;


    public DorisStreamLoad(String hostPort, String db, String tbl, String user, String passwd) {
        this.hostPort = hostPort;
        this.db = db;
        this.tbl = tbl;
        this.user = user;
        this.passwd = passwd;
        this.authEncoding = Base64.getEncoder().encodeToString(String.format("%s:%s", user, StringUtils.isBlank(passwd) ? "" : passwd).getBytes(StandardCharsets.UTF_8));
    }

    private String getLoadUrlStr() {
        return String.format(loadUrlPattern, hostPort, db, tbl);
    }

    //获取http连接信息
    private HttpURLConnection getConnection(String urlStr, String label, String columns, String jsonformat) throws IOException {
        URL url = new URL(urlStr);
        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setInstanceFollowRedirects(false);
        conn.setRequestMethod("PUT");
        conn.setRequestProperty("Authorization", "Basic " + authEncoding);
        conn.addRequestProperty("Expect", "100-continue");
        conn.setRequestProperty("Content-Type", "text/html; charset=UTF-8");

        conn.setConnectTimeout(30000);
        conn.setReadTimeout(30000);
        conn.setUseCaches(false);
        conn.setRequestProperty("Connection", "Keep-Alive");

        conn.addRequestProperty("label", label);
        conn.addRequestProperty("max_filter_ratio", "0");
        conn.addRequestProperty("strict_mode", "true");
        if (columns != null && columns.trim().length() != 0)
            conn.addRequestProperty("columns", columns);
        conn.addRequestProperty("format", "json");
        if (jsonformat != null && jsonformat.trim().length() != 0)
            conn.addRequestProperty("jsonpaths", jsonformat);
        conn.addRequestProperty("strip_outer_array", "true");
        conn.setDoOutput(true);
        conn.setDoInput(true);

        return conn;
    }

    //执行数据导入
    public LoadResponse loadBatchWithHttpClient(String data,
                                                List<String> columns,
                                                String jsonformat,
                                                Map<Object, Object> redirectMapping) throws IOException {
        Calendar calendar = Calendar.getInstance();
        //导入的lable，全局唯一
        String label = String.format("flink_conn-%s_%s-%s%02d%02d_%02d%02d%02d-%s", this.db, this.tbl,
                calendar.get(Calendar.YEAR), calendar.get(Calendar.MONTH) + 1, calendar.get(Calendar.DAY_OF_MONTH),
                calendar.get(Calendar.HOUR_OF_DAY), calendar.get(Calendar.MINUTE), calendar.get(Calendar.SECOND),
                UUID.randomUUID().toString().replaceAll("-", "_"));

        long start = System.currentTimeMillis();
        final String host = hostPort.split("\\.")[0];
        if (null == host) {
            throw new RuntimeException("None of the load url can be connected.");
        }
        // do http put request and get response
        LoadResponse loadResponse = this.doHttpPut(data, columns, label, jsonformat, redirectMapping);

        long cost = System.currentTimeMillis() - start;
        log.trace("StreamLoad response: " + JSON.toJSONString(loadResponse) + ", cost(ms): " + cost);
        final String keyStatus = "Status";
        if (null == loadResponse || loadResponse.getStatus() == null) {
            throw new RuntimeException("Unable to flush data to doris: unknown result status. resp result:" + JSON.toJSONString(loadResponse));
        }
        return loadResponse;
    }

    private LoadResponse doHttpPut(String data, List<String> columns, String label, String jsonformat, Map<Object, Object> redirectMapping) throws IOException {
        HttpClientBuilder httpClientBuilder = HttpClients.custom()
                //TODO :y
                .addInterceptorFirst((HttpRequestInterceptor) (httpRequest, httpContext) -> httpRequest.removeHeaders(HttpHeaders.CONTENT_LENGTH))
                .setRedirectStrategy(new DefaultRedirectStrategy() {
                    @Override
                    protected boolean isRedirectable(final String method) {
                        return true;
                    }

                    @Override
                    public HttpUriRequest getRedirect(HttpRequest request, HttpResponse response, HttpContext context) throws ProtocolException {
                        URI uri = this.getLocationURI(request, response, context);
                        //reset redirect url
                        if (redirectMapping != null) {
                            String finalLocation = uri.toString();
                            Optional<Map.Entry<Object, Object>> firstOpt = redirectMapping.entrySet().stream().filter(x -> finalLocation.equals(x.getKey().toString())).findFirst();
                            if (firstOpt.isPresent()) {
                                try {
                                    uri = new URI(firstOpt.get().getValue().toString());
                                } catch (URISyntaxException e) {
                                    e.printStackTrace();
                                }
                            }
                        }
                        String method = request.getRequestLine().getMethod();
                        if (method.equalsIgnoreCase("HEAD")) {
                            return new HttpHead(uri);
                        } else if (method.equalsIgnoreCase("GET")) {
                            return new HttpGet(uri);
                        } else {
                            int status = response.getStatusLine().getStatusCode();
                            return status == 307 ? RequestBuilder.copy(request).setUri(uri).build() : new HttpGet(uri);
                        }
                    }
                });

        try (CloseableHttpClient httpclient = httpClientBuilder.build()) {
            HttpPut httpPut = new HttpPut(getLoadUrlStr());
            List<String> cols = columns;
            if (null != cols && !cols.isEmpty()) {
                httpPut.setHeader("columns", String.join(",", cols));
            }
            if (jsonformat != null && jsonformat.trim().length() != 0)
                httpPut.setHeader("jsonpaths", jsonformat);

            // set other required headers
            httpPut.setHeader(HttpHeaders.EXPECT, "100-continue");
            httpPut.setHeader(HttpHeaders.AUTHORIZATION, "Basic " + authEncoding);
            httpPut.setHeader("label", label);
            httpPut.setHeader("max_filter_ratio", "0");
            httpPut.setHeader("strict_mode", "true");
            httpPut.setHeader("format", "json");
            httpPut.setHeader("strip_outer_array", "true");

            // Use ByteArrayEntity instead of StringEntity to handle Chinese correctly
            httpPut.setEntity(new StringEntity(data));

            try (CloseableHttpResponse resp = httpclient.execute(httpPut)) {
                final int code = resp.getStatusLine().getStatusCode();
                if (HttpStatus.SC_OK != code) {
                    log.warn("Request failed with code:{}, response body:{}", code, EntityUtils.toString(resp.getEntity()));
                    return null;
                }
                HttpEntity respEntity = resp.getEntity();
                if (null == respEntity) {
                    log.warn("Request failed with empty response.");
                    return null;
                }

                return new LoadResponse(code, String.valueOf(HttpStatus.SC_OK != code ? HttpStatus.SC_INTERNAL_SERVER_ERROR : HttpStatus.SC_OK), JSON.parseObject(EntityUtils.toString(respEntity), LoadResponseContext.class));
            } catch (Exception e) {
                e.printStackTrace();
                throw e;
            }
        }
    }

    //执行数据导入
    public LoadResponse loadBatch(String JsonData, String columns, String jsonformat, Map<Object, Object> redirectMapping) {
        Calendar calendar = Calendar.getInstance();
        //导入的lable，全局唯一
        String label = String.format("flink_conn-%s_%s-%s%02d%02d_%02d%02d%02d-%s", this.db, this.tbl,
                calendar.get(Calendar.YEAR), calendar.get(Calendar.MONTH) + 1, calendar.get(Calendar.DAY_OF_MONTH),
                calendar.get(Calendar.HOUR_OF_DAY), calendar.get(Calendar.MINUTE), calendar.get(Calendar.SECOND),
                UUID.randomUUID().toString().replaceAll("-", "_"));

        HttpURLConnection feConn = null;
        HttpURLConnection beConn = null;
        StringBuilder response = new StringBuilder();
        try {
            // build request and send to fe
            feConn = getConnection(getLoadUrlStr(), label, columns, jsonformat);
            int status = feConn.getResponseCode();
            // get fe respond
            InputStream stream = null;
            BufferedReader br = null;
            String line;
            try {
                stream = (InputStream) feConn.getContent();
                br = new BufferedReader(new InputStreamReader(stream));
                while ((line = br.readLine()) != null) {
                    response.append(line);
                }
            } finally {
                if (stream != null)
                    stream.close();
                if (br != null)
                    br.close();
            }

            // fe send back http response code TEMPORARY_REDIRECT 307 and new be location
            if (status != 307) {
                throw new Exception("status is not TEMPORARY_REDIRECT 307, status: " + status + " feResp:" + response + " url:" + getLoadUrlStr());
            }
            String location = feConn.getHeaderField("Location");
            if (location == null) {
                throw new Exception("redirect location is null. feResp:" + response + " url:" + getLoadUrlStr());
            }

            //reset redirect url
            if (redirectMapping != null) {
                final String finalLocation = location;
                Optional<Map.Entry<Object, Object>> firstOpt = redirectMapping.entrySet().stream().filter(x -> finalLocation.equals(x.getKey().toString())).findFirst();
                if (firstOpt.isPresent())
                    location = firstOpt.get().getValue().toString();
            }

            // build request and send to new be location
            beConn = getConnection(location, label, columns, jsonformat);
            // send data to be
            BufferedOutputStream bos = new BufferedOutputStream(beConn.getOutputStream());
            bos.write(JsonData.getBytes());
            bos.close();

            // get respond
            try {
                line = null;
                status = beConn.getResponseCode();
                String respMsg = beConn.getResponseMessage();
                stream = (InputStream) beConn.getContent();
                br = new BufferedReader(new InputStreamReader(stream));
                response = new StringBuilder();
                while ((line = br.readLine()) != null) {
                    response.append(line);
                }
                return new LoadResponse(status, respMsg, JSON.parseObject(response.toString(), LoadResponseContext.class));
            } finally {
                if (stream != null)
                    stream.close();
                if (br != null)
                    br.close();
            }

        } catch (Exception e) {
            e.printStackTrace();
            String err = "failed to load audit via AuditLoader plugin with label: " + label + " errMsg:" + e.getMessage() + " url:" + getLoadUrlStr() + " resp:" + response;
            log.warn(err);
            return new LoadResponse(-1, err, null);
        } finally {
            if (feConn != null) {
                feConn.disconnect();
            }
            if (beConn != null) {
                beConn.disconnect();
            }
        }
    }


    public String getDb() {
        return db;
    }

    public void setDb(String db) {
        this.db = db;
    }

    public String getTbl() {
        return tbl;
    }

    public void setTbl(String tbl) {
        this.tbl = tbl;
    }

    public static class LoadResponse {
        private Integer status;
        private String respMsg;
        private LoadResponseContext respContent;

        public LoadResponse(int status, String respMsg, LoadResponseContext respContent) {
            this.status = status;
            this.respMsg = respMsg;
            this.respContent = respContent;
        }

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder();
            sb.append("status: ").append(status);
            sb.append(", resp msg: ").append(respMsg);
            sb.append(", resp content: ").append(respContent);
            return sb.toString();
        }

        public Integer getStatus() {
            return status;
        }

        public void setStatus(Integer status) {
            this.status = status;
        }

        public String getRespMsg() {
            return respMsg;
        }

        public void setRespMsg(String respMsg) {
            this.respMsg = respMsg;
        }

        public LoadResponseContext getRespContent() {
            return respContent;
        }

        public void setRespContent(LoadResponseContext respContent) {
            this.respContent = respContent;
        }

        public boolean isOk() {
            if (this != null) {
                Integer status = this.getStatus();
                if (status != null && status == 200) {
                    LoadResponseContext respContent = this.getRespContent();
                    List<String> successStatus = Arrays.asList("Success", "Publish Timeout");
                    if (successStatus.contains(respContent.getStatus()) && respContent.getNumberTotalRows() == respContent.getNumberLoadedRows())
                        return true;
                }
            }
            return false;
        }
    }

    /**
     * TxnId：导入的事务ID。用户可不感知。
     * <p>
     * Status: 导入最后的状态。
     * <p>
     * Success：表示导入成功，数据已经可见。
     * <p>
     * Publish Timeout：表述导入作业已经成功Commit，但是由于某种原因并不能立即可见。用户可以视作已经成功不必重试导入。
     * <p>
     * Label Already Exists：表明该Label已经被其他作业占用，可能是导入成功，也可能是正在导入。
     * <p>
     * 其他：此次导入失败，用户可以指定Label重试此次作业。
     * <p>
     * Message: 导入状态的详细说明。失败时会返回具体的失败原因。
     * <p>
     * NumberTotalRows: 从数据流中读取到的总行数。
     * <p>
     * NumberLoadedRows: 此次导入的数据行数，只有在Success时有效。
     * <p>
     * NumberFilteredRows: 此次导入过滤掉的行数，即数据质量不合格的行。
     * <p>
     * NumberUnselectedRows: 此次导入，通过 where 条件被过滤掉的行数。
     * <p>
     * LoadBytes: 此次导入的源文件数据量大小。
     * <p>
     * LoadTimeMs: 此次导入所用的时间(ms)。
     * <p>
     * ErrorURL: 被过滤数据的具体内容，仅保留前1000条。如果导入任务失败，可以直接用以下方式获取被过滤的数据，并进行分析，以调整导入任务。
     */
    public static class LoadResponseContext {
        private Long TxnId;
        private String Label;
        private String Status;
        private String Message;
        private Long NumberTotalRows;
        private Long NumberLoadedRows;
        private Long NumberFilteredRows;
        private Long NumberUnselectedRows;
        private Long LoadBytes;
        private Long LoadTimeMs;
        private Long BeginTxnTimeMs;
        private Long StreamLoadPutTimeMs;
        private Long ReadDataTimeMs;
        private Long WriteDataTimeMs;
        private Long CommitAndPublishTimeMs;

        public Long getTxnId() {
            return TxnId;
        }

        public void setTxnId(Long txnId) {
            TxnId = txnId;
        }

        public String getLabel() {
            return Label;
        }

        public void setLabel(String label) {
            Label = label;
        }

        public String getStatus() {
            return Status;
        }

        public void setStatus(String status) {
            Status = status;
        }

        public String getMessage() {
            return Message;
        }

        public void setMessage(String message) {
            Message = message;
        }

        public Long getNumberTotalRows() {
            return NumberTotalRows;
        }

        public void setNumberTotalRows(Long numberTotalRows) {
            NumberTotalRows = numberTotalRows;
        }

        public Long getNumberLoadedRows() {
            return NumberLoadedRows;
        }

        public void setNumberLoadedRows(Long numberLoadedRows) {
            NumberLoadedRows = numberLoadedRows;
        }

        public Long getNumberFilteredRows() {
            return NumberFilteredRows;
        }

        public void setNumberFilteredRows(Long numberFilteredRows) {
            NumberFilteredRows = numberFilteredRows;
        }

        public Long getNumberUnselectedRows() {
            return NumberUnselectedRows;
        }

        public void setNumberUnselectedRows(Long numberUnselectedRows) {
            NumberUnselectedRows = numberUnselectedRows;
        }

        public Long getLoadBytes() {
            return LoadBytes;
        }

        public void setLoadBytes(Long loadBytes) {
            LoadBytes = loadBytes;
        }

        public Long getLoadTimeMs() {
            return LoadTimeMs;
        }

        public void setLoadTimeMs(Long loadTimeMs) {
            LoadTimeMs = loadTimeMs;
        }

        public Long getBeginTxnTimeMs() {
            return BeginTxnTimeMs;
        }

        public void setBeginTxnTimeMs(Long beginTxnTimeMs) {
            BeginTxnTimeMs = beginTxnTimeMs;
        }

        public Long getStreamLoadPutTimeMs() {
            return StreamLoadPutTimeMs;
        }

        public void setStreamLoadPutTimeMs(Long streamLoadPutTimeMs) {
            StreamLoadPutTimeMs = streamLoadPutTimeMs;
        }

        public Long getReadDataTimeMs() {
            return ReadDataTimeMs;
        }

        public void setReadDataTimeMs(Long readDataTimeMs) {
            ReadDataTimeMs = readDataTimeMs;
        }

        public Long getWriteDataTimeMs() {
            return WriteDataTimeMs;
        }

        public void setWriteDataTimeMs(Long writeDataTimeMs) {
            WriteDataTimeMs = writeDataTimeMs;
        }

        public Long getCommitAndPublishTimeMs() {
            return CommitAndPublishTimeMs;
        }

        public void setCommitAndPublishTimeMs(Long commitAndPublishTimeMs) {
            CommitAndPublishTimeMs = commitAndPublishTimeMs;
        }

        @Override
        public String toString() {
            return "LoadResponseContext{" +
                    "TxnId=" + TxnId +
                    ", Label='" + Label + '\'' +
                    ", Status='" + Status + '\'' +
                    ", Message='" + Message + '\'' +
                    ", NumberTotalRows=" + NumberTotalRows +
                    ", NumberLoadedRows=" + NumberLoadedRows +
                    ", NumberFilteredRows=" + NumberFilteredRows +
                    ", NumberUnselectedRows=" + NumberUnselectedRows +
                    ", LoadBytes=" + LoadBytes +
                    ", LoadTimeMs=" + LoadTimeMs +
                    ", BeginTxnTimeMs=" + BeginTxnTimeMs +
                    ", StreamLoadPutTimeMs=" + StreamLoadPutTimeMs +
                    ", ReadDataTimeMs=" + ReadDataTimeMs +
                    ", WriteDataTimeMs=" + WriteDataTimeMs +
                    ", CommitAndPublishTimeMs=" + CommitAndPublishTimeMs +
                    '}';
        }
    }


}
